package ai.mspbots.poc.controller.v2;

import ai.mspbots.poc.controller.v2.dto.IncrementDTO;
import ai.mspbots.poc.controller.v2.dto.PageVO;
import ai.mspbots.poc.es.entity.Document;
import ai.mspbots.poc.help.PocEnv;
import ai.mspbots.poc.help.SpellCheckHelper;
import ai.mspbots.poc.help.VectorHelper;
import cn.hutool.core.util.StrUtil;
import cn.hutool.extra.spring.SpringUtil;
import cn.hutool.json.JSONUtil;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.common.lucene.search.function.FunctionScoreQuery;
import org.elasticsearch.common.unit.Fuzziness;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.MatchPhraseQueryBuilder;
import org.elasticsearch.index.query.MatchQueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.functionscore.FunctionScoreQueryBuilder;
import org.elasticsearch.index.query.functionscore.ScriptScoreFunctionBuilder;
import org.elasticsearch.script.Script;
import org.elasticsearch.script.ScriptType;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Component
public class SearchConfig {

    private static final String INDEX = SpringUtil.getBean(PocEnv.class).getIndex();
    private final RestHighLevelClient client;

    public SearchConfig(RestHighLevelClient client) {
        this.client = client;
    }

    public PageVO<Document> search(String searchText, float[] vector, int page, int size) throws Exception {
        SearchRequest searchRequest = new SearchRequest(INDEX);

        BoolQueryBuilder boolQuery = new BoolQueryBuilder();

        String correctedText = SpellCheckHelper.correctGrammar(searchText);

        // must 子句
        boolQuery.should(new MatchPhraseQueryBuilder("name", searchText).boost(10.0f));
        boolQuery.should(new MatchPhraseQueryBuilder("description", searchText).boost(8.0f));
        boolQuery.should(new MatchPhraseQueryBuilder("name", correctedText).boost(10.0f));
        boolQuery.should(new MatchPhraseQueryBuilder("description", correctedText).boost(8.0f));

        // should 子句
        boolQuery.should(new MatchQueryBuilder("description", searchText).fuzziness(Fuzziness.fromEdits(2)).fuzzyTranspositions(false).boost(7.0f));
        boolQuery.should(new MatchQueryBuilder("name", searchText).fuzziness(Fuzziness.fromEdits(2)).fuzzyTranspositions(false).boost(7.0f));

        // 添加 name_vector 的相似度查询
        Map<String, Object> nameParams = new HashMap<>();
        nameParams.put("query_vector", vector);
        // 计算欧几里得距离
        Script nameVectorScript =
                new Script(
                        ScriptType.INLINE,
                        "painless",
                        """
                        double dotProduct = 0.0;
                         double normA = 0.0;
                         double normB = 0.0;
                         if (doc['name_vector'] == null || doc['name_vector'].size() == 0) {
                             return 0;
                         }
             
                         List userVector = params.query_vector;
             
                         for (int i = 0; i < userVector.size(); i++) {
                           dotProduct += userVector.get(i) * doc['name_vector'].vectorValue[i];
                           normA += Math.pow(userVector.get(i), 2);
                           normB += Math.pow(doc['name_vector'].vectorValue[i], 2);
                         }
             
                         normA = Math.sqrt(normA);
                         normB = Math.sqrt(normB);
                         if (normA == 0 || normB == 0) {
                           return 0;
                         }
                         double score = dotProduct / (normA * normB);
                         return Math.max(0, score);
                                """,
                        nameParams
                );

        ScriptScoreFunctionBuilder nameVectorScriptScoreFunction = new ScriptScoreFunctionBuilder(nameVectorScript);
        nameVectorScriptScoreFunction.setWeight(3.0f);
        // 添加 description_vector 的相似度查询
        Map<String, Object> descriptionParams = new HashMap<>();
        descriptionParams.put("query_vector", vector);
        // 计算欧几里得距离
        Script descriptionVectorScript = new Script(
                ScriptType.INLINE,
                "painless",
                """
                     double dotProduct = 0.0;
                     double normA = 0.0;
                     double normB = 0.0;
                     if (doc['name_vector'] == null || doc['name_vector'].size() == 0) {
                         return 0;
                     }
         
                     List userVector = params.query_vector;
         
                     for (int i = 0; i < userVector.size(); i++) {
                       dotProduct += userVector.get(i) * doc['name_vector'].vectorValue[i];
                       normA += Math.pow(userVector.get(i), 2);
                       normB += Math.pow(doc['name_vector'].vectorValue[i], 2);
                     }
         
                     normA = Math.sqrt(normA);
                     normB = Math.sqrt(normB);
                     if (normA == 0 || normB == 0) {
                       return 0;
                     }
                     double score = dotProduct / (normA * normB);
                     return Math.max(0, score);
                        """,
                descriptionParams
        );

        ScriptScoreFunctionBuilder descriptionVectorScriptScoreFunction = new ScriptScoreFunctionBuilder(descriptionVectorScript);
        descriptionVectorScriptScoreFunction.setWeight(2.0f);

        // 查询权重
        Script weightScript = new Script(
                ScriptType.INLINE,
                "painless",
                """
                          if (doc['weight'].size() == 0) {
                            return 0;
                          } else {
                            return doc['weight'].value * params.userWeight;
                          }
                        """,
                Map.of("userWeight", 1)
        );
        ScriptScoreFunctionBuilder weightScriptScoreFunction = new ScriptScoreFunctionBuilder(weightScript);
        weightScriptScoreFunction.setWeight(1.0f);

        // 将脚本评分查询添加到布尔查询中
        FunctionScoreQueryBuilder functionScoreQuery = new FunctionScoreQueryBuilder(
                boolQuery,
                new FunctionScoreQueryBuilder.FilterFunctionBuilder[]{
                        new FunctionScoreQueryBuilder.FilterFunctionBuilder(nameVectorScriptScoreFunction),
                        new FunctionScoreQueryBuilder.FilterFunctionBuilder(descriptionVectorScriptScoreFunction),
                        new FunctionScoreQueryBuilder.FilterFunctionBuilder(weightScriptScoreFunction)
                }
        ).scoreMode(FunctionScoreQuery.ScoreMode.SUM).boost(100.0f); // 或者选择其他模式
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        // 设置分页
        searchSourceBuilder.from((page - 1) * size); // 设置起始位置
        searchSourceBuilder.size(size); // 设置每页返回的文档数量
        // 将查询添加到搜索源
        searchSourceBuilder.query(functionScoreQuery);
        searchRequest.source(searchSourceBuilder);
        // 执行查询
        SearchResponse searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
        List<Document> documents = new ArrayList<>();
        for (SearchHit hit : searchResponse.getHits().getHits()) {
            String source = hit.getSourceAsString();
            Document doc = JSONUtil.toBean(source, Document.class);
            float score = hit.getScore();
            doc.setScore(score);
            // 继续提取其他字段
            documents.add(doc);
        }
        long total = searchResponse.getHits().getTotalHits().value;

        PageVO<Document> pageVO = new PageVO<>();
        pageVO.setTotal((int) total);
        pageVO.setRecords(documents);
        pageVO.setPage(page);
        pageVO.setSize(size);
        return pageVO;
    }


    public PageVO<Document> search(String searchText, int page, int size) throws Exception {
        SearchRequest searchRequest = new SearchRequest(INDEX);
        BoolQueryBuilder boolQuery = new BoolQueryBuilder();

        String correctedText = SpellCheckHelper.correctGrammar(searchText);

        // should 子句
        boolQuery.should(new MatchPhraseQueryBuilder("name", searchText).boost(10.0f));
        boolQuery.should(new MatchPhraseQueryBuilder("description", searchText).boost(8.0f));
        boolQuery.should(new MatchPhraseQueryBuilder("name", correctedText).boost(10.0f));
        boolQuery.should(new MatchPhraseQueryBuilder("description", correctedText).boost(8.0f));
        boolQuery.should(new MatchQueryBuilder("description", searchText).fuzziness(Fuzziness.AUTO).fuzzyTranspositions(false).prefixLength(2).boost(7.0f));
        boolQuery.should(new MatchQueryBuilder("name", searchText).fuzziness(Fuzziness.AUTO).fuzzyTranspositions(false).prefixLength(2).boost(0.7f));

        // 创建搜索源并设置分页
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        searchSourceBuilder.from((page - 1) * size); // 设置起始位置
        searchSourceBuilder.size(size); // 设置每页返回的文档数量
        // 将查询添加到搜索源
        searchSourceBuilder.query(boolQuery);
        searchRequest.source(searchSourceBuilder);
        // 执行查询
        SearchResponse searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
        List<Document> documents = new ArrayList<>();
        for (SearchHit hit : searchResponse.getHits().getHits()) {
            String source = hit.getSourceAsString();
            Document doc = JSONUtil.toBean(source, Document.class);
            float score = hit.getScore();
            doc.setScore(score);
            // 继续提取其他字段
            documents.add(doc);
        }
        long total = searchResponse.getHits().getTotalHits().value;

        PageVO<Document> pageVO = new PageVO<>();
        pageVO.setTotal((int) total);
        pageVO.setRecords(documents);
        pageVO.setPage(page);
        pageVO.setSize(size);
        return pageVO;
    }

    public INDArray vectorized(String text) {
        if (StrUtil.isEmpty(text)) {
            return null;
        }
        return VectorHelper.vectorizedText(text);
    }

    public void updateDocumentWeight(IncrementDTO incrementDTO) throws Exception {
        UpdateRequest updateRequest = new UpdateRequest(INDEX, incrementDTO.getDocumentId());
        // 查询权重
        Script weightScript = new Script(
                ScriptType.INLINE,
                "painless",
                """
                          if (ctx._source.weight == null) { ctx._source.weight = 0 } ctx._source.weight += params.increment
                        """,
                Map.of("increment", incrementDTO.getIncrement())
        );
        updateRequest.script(weightScript);
        client.update(updateRequest, RequestOptions.DEFAULT);
    }

    public SearchResponse searchWithWeight(String index, double userWeight) throws Exception {
        SearchRequest searchRequest = new SearchRequest(index);
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();

        Map<String, Object> descriptionParams = new HashMap<>();
        descriptionParams.put("userWeight", userWeight);
        // 查询权重
        Script weightScript = new Script(
                ScriptType.INLINE,
                "painless",
                """
                          doc['weight'].value * params.userWeight
                        """,
                descriptionParams
        );
        ScriptScoreFunctionBuilder weightScriptScoreFunction = new ScriptScoreFunctionBuilder(weightScript);
        weightScriptScoreFunction.setWeight(2.0f);
        // 使用 function_score 来调整得分
        sourceBuilder.query(QueryBuilders.functionScoreQuery(
                QueryBuilders.matchAllQuery(), // 你的查询条件
                new FunctionScoreQueryBuilder.FilterFunctionBuilder[]{
                        new FunctionScoreQueryBuilder.FilterFunctionBuilder(
                                weightScriptScoreFunction
                        )
                }
        ).scoreMode(FunctionScoreQuery.ScoreMode.SUM));
        searchRequest.source(sourceBuilder);
        return client.search(searchRequest, RequestOptions.DEFAULT);
    }
}
